# -*-coding:utf-8-*-
import cv2
import numpy as np
import random


class Roughness:
    def __init__(self, img):
        if isinstance(img, str):
            self.img = cv2.imread(img, 0)
        else:
            self.img = img
        self.Sa = 0
        self.forward()

    # 最小二乘拟合
    def LS_fitting(self, xx, yy, Z) -> np.array:
        """
        最小二乘拟合，处理中产生拟合平面，并返回原始平面-拟合平面，调平处理
        :param X: 平面X坐标
        :param Y: 平面Y坐标
        :param Z: 调平之后平面Z
        :return: Z
        """
        X, Y = np.meshgrid(xx, yy)
        x1 = random.choices(xx, k=100)
        x2 = random.choices(yy, k=100)
        # 创建系数矩阵A
        A = np.zeros((3, 3))
        for i in range(0, 100):
            A[0, 0] = A[0, 0] + x1[i] ** 2
            A[0, 1] = A[0, 1] + x1[i] * x2[i]
            A[0, 2] = A[0, 2] + x1[i]
            A[1, 0] = A[0, 1]
            A[1, 1] = A[1, 1] + x2[i] ** 2
            A[1, 2] = A[1, 2] + x2[i]
            A[2, 0] = A[0, 2]
            A[2, 1] = A[1, 2]
            A[2, 2] = 100

        # 创建b
        b = np.zeros((3, 1))
        for i in range(0, 100):
            for a in range(640):
                if xx[a] == x1[i]:
                    y_index = a
            for c in range(480):
                if yy[c] == x2[i]:
                    x_index = c
            b[0, 0] = b[0, 0] + x1[i] * Z[x_index][y_index]
            b[1, 0] = b[1, 0] + x2[i] * Z[x_index][y_index]
            b[2, 0] = b[2, 0] + Z[x_index][y_index]

        # 求解X
        A_inv = np.linalg.inv(A)
        x = np.dot(A_inv, b)
        z = Z - (x[0, 0] * X + x[1, 0] * Y + x[2, 0])
        return z

    # 计算粗糙度Sa
    def calculate_Sa(self, Z_correct):
        """
        计算粗糙度Sa
        """
        Sa = np.mean(np.abs(Z_correct))
        return Sa

    @staticmethod
    def mean_val(img):
        total = 0
        num = 0
        for i in range(img.shape[0]):
            for j in range(img.shape[1]):
                if img[i, j] != 0:
                    total += img[i, j]
                    num += 1
        return np.round(total / num, 3)

    @staticmethod
    def fixup(img):
        val = Roughness.mean_val(img)
        # print(val)
        for i in range(img.shape[0]):
            for j in range(img.shape[1]):
                if img[i, j] == 0:
                    img[i, j] = val
        return img

    def forward(self):
        """
        初始化自动调用，直接取Sa即可
        """
        if self.img.shape[0] == 480:
            fix_up = Roughness.fixup(self.img)
            Z_correct = np.round(self.LS_fitting(Roughness.xx, Roughness.yy, fix_up), 3)
            Sa = self.calculate_Sa(Z_correct)
        else:
            Sa = np.mean(np.abs(self.img - np.mean(self.img)))
        self.Sa = np.round(Sa, 3)

    xx = np.linspace(0, 315, 640)
    yy = np.linspace(0, 236, 480)


if __name__ == '__main__':
    Sa = Roughness("../images/test.png")
    print(Sa.Sa)






